import os
import shutil
import subprocess

# 定义ROS bag文件所在的目录和激光雷达点云话题
bag_dir = '/home/cc/Workspace/ros_ws/src/lidar_camera_calibration/data/lidar'
lidar_topic = '/unilidar/cloud'

# 获取目录下所有的ROS bag文件
bag_files = [f for f in os.listdir(bag_dir) if f.endswith('.bag')]

for bag_file in bag_files:
    bag_path = os.path.join(bag_dir, bag_file)
    # 创建一个新的文件夹来保存当前bag文件转换后的PCD文件
    pcd_dir = os.path.splitext(bag_path)[0]
    # 检查文件夹是否已经存在，如果存在则删除
    if os.path.exists(pcd_dir):
        shutil.rmtree(pcd_dir)
    os.makedirs(pcd_dir, exist_ok=False)
    # 播放ROS bag文件
    play_process = subprocess.Popen(['rosbag', 'play', bag_path])
    try:
        # 订阅话题并保存为PCD文件
        pcd_process = subprocess.Popen(['rosrun', 'pcl_ros', 'pointcloud_to_pcd', f'input:={lidar_topic}'], cwd=pcd_dir)
        # 等待rosbag播放完成
        play_process.wait()
        # 终止pcl_ros节点
        pcd_process.terminate()

        # 获取当前文件夹中的所有PCD文件
        pcd_files = [os.path.join(pcd_dir, f) for f in os.listdir(pcd_dir) if f.endswith('.pcd')]
        if pcd_files:
            # 合并当前文件夹中的所有PCD文件
            try:
                merge_command = ['pcl_concatenate_points_pcd'] + pcd_files
                subprocess.run(merge_command, cwd=pcd_dir, check=True)
                print(f"Successfully merged PCD files in {pcd_dir}")

                # 假设合并后的文件名为 concatenated_cloud.pcd
                merged_file = os.path.join(pcd_dir, 'output.pcd')
                if os.path.exists(merged_file):
                    # 删除除合并后文件之外的其他PCD文件
                    for file in pcd_files:
                        if file != merged_file:
                            os.remove(file)
                    print(f"Deleted other PCD files in {pcd_dir}, only kept the merged file.")
                else:
                    print(f"Could not find the merged PCD file in {pcd_dir}")
            except subprocess.CalledProcessError as e:
                print(f"Error merging PCD files in {pcd_dir}: {e}")
        else:
            print(f"No PCD files found in {pcd_dir}")

    except KeyboardInterrupt:
        # 手动停止时，终止所有进程
        play_process.terminate()
        pcd_process.terminate()

# import os
# import shutil
# import subprocess

# # 定义ROS bag文件所在的目录和激光雷达点云话题
# bag_dir = '/home/cc/Workspace/ros_ws/src/lidar_camera_calibration/data/lidar'
# lidar_topic = '/unilidar/cloud'

# # 获取目录下所有的ROS bag文件
# bag_files = [f for f in os.listdir(bag_dir) if f.endswith('.bag')]

# for bag_file in bag_files:
#     bag_path = os.path.join(bag_dir, bag_file)
#     # 创建一个新的文件夹来保存当前bag文件转换后的PCD文件
#     pcd_dir = os.path.splitext(bag_path)[0]
#     # 检查文件夹是否已经存在，如果存在则删除
#     if os.path.exists(pcd_dir):
#         shutil.rmtree(pcd_dir)
#     os.makedirs(pcd_dir, exist_ok=False)
#     # 播放ROS bag文件
#     play_process = subprocess.Popen(['rosbag', 'play', bag_path])
#     try:
#         # 订阅话题并保存为PCD文件
#         pcd_process = subprocess.Popen(['rosrun', 'pcl_ros', 'pointcloud_to_pcd', f'input:={lidar_topic}'], cwd=pcd_dir)
#         # 等待rosbag播放完成
#         play_process.wait()
#         # 终止pcl_ros节点
#         pcd_process.terminate()

#     except KeyboardInterrupt:
#         # 手动停止时，终止所有进程
#         play_process.terminate()
#         pcd_process.terminate()
